#include "MyWindow.hpp"
#include <algorithm>
#include <assert.h>

static bool g_classRegistered = false;
static const wchar_t* const WINDOW_CLASS_NAME = L"MyWindowClass";

void Window::InitializeWindow()
{
    if (!g_classRegistered) {
        RegisterWindowClass();
        g_classRegistered = true;
    }
}

void Window::DestroyWindow()
{
    if (g_classRegistered) {
        UnregisterWindowClass();
        g_classRegistered = false;
    }
}

void Window::RegisterWindowClass()
{
    ATOM class_atom;

    WNDCLASSW window_class{};   

    window_class.style         = CS_HREDRAW | CS_VREDRAW;
    window_class.lpfnWndProc   = Window::WindowProc;
    window_class.cbClsExtra    = 0;
    window_class.cbWndExtra    = 0;
    window_class.hInstance     = GetModuleHandleW(NULL);

    window_class.hIcon         = LoadIcon(NULL, IDI_APPLICATION); 
    window_class.hCursor       = LoadCursor(NULL, IDC_ARROW);     
    window_class.hbrBackground = (HBRUSH)GetStockObject(WHITE_BRUSH);
    window_class.lpszMenuName  = NULL;
    window_class.lpszClassName = WINDOW_CLASS_NAME;             

    class_atom = RegisterClassW(&window_class); 

    assert(class_atom != 0);
    (void)class_atom; 
}

void Window::UnregisterWindowClass()
{
    UnregisterClassW(WINDOW_CLASS_NAME, GetModuleHandleW(nullptr));
}

Window::Window(const char* title, int width, int height)
{
    int len = MultiByteToWideChar(CP_UTF8, 0, title, -1, nullptr, 0);
    std::wstring wtitle(len, 0);
    MultiByteToWideChar(CP_UTF8, 0, title, -1, &wtitle[0], len);

    m_handle = CreateWindowExW(
        0,
        WINDOW_CLASS_NAME,
        wtitle.c_str(),
        WS_OVERLAPPEDWINDOW,
        CW_USEDEFAULT, CW_USEDEFAULT,
        width,
        height,
        nullptr,
        nullptr,
        GetModuleHandle(nullptr),
        this                  
    );

    if (!m_handle)
        throw std::runtime_error("CreateWindowExW failed");

    ShowWindow(m_handle, SW_SHOW);
    UpdateWindow(m_handle);

    HDC hdc = GetDC(m_handle);

    m_memoryDC = CreateCompatibleDC(hdc);

    CreateSurface(width, height);

    ReleaseDC(m_handle, hdc);
}

Window::~Window()
{
    if (m_hBitmap)
        DeleteObject(m_hBitmap);

    if (m_memoryDC)
        DeleteDC(m_memoryDC);

    if (m_surface)
        delete m_surface;

    if (m_handle)
        ::DestroyWindow(m_handle); 
}

void Window::PollEvents()
{
    MSG msg;
    while (PeekMessage(&msg, nullptr, 0, 0, PM_REMOVE))
    {
        if (msg.message == WM_QUIT)
        {
            m_shouldClose = true;
            return;
        }

        TranslateMessage(&msg);
        DispatchMessage(&msg);
    }
}

bool Window::IsKeyPressed(KeyCode key) const
{
    return m_keyDown[(int)key];
}

bool Window::IsMousePressed(MouseButton button) const
{
    return m_mouseDown[(int)button];
}

void Window::GetCursorPos(float &x, float &y) const
{
    POINT p;
    ::GetCursorPos(&p);               
    ::ScreenToClient(m_handle, &p);   

    x = (float)p.x;
    y = (float)p.y;
}

void Window::HandleKeyMessage(WPARAM virtualKey, bool pressed)
{
    KeyCode key = KeyCode::Count;

    switch (virtualKey)
    {
        case 'A':      key = KeyCode::A;     break;
        case 'D':      key = KeyCode::D;     break;
        case 'S':      key = KeyCode::S;     break;
        case 'W':      key = KeyCode::W;     break;
        case VK_SPACE: key = KeyCode::Space; break;
    }

    if (key != KeyCode::Count)
    {
        m_keyDown[(int)key] = pressed;

        if (m_callbacks.onKey)
            m_callbacks.onKey(key, pressed);
    }
}

void Window::HandleMouseButtonMessage(MouseButton button, bool pressed)
{
    m_mouseDown[(int)button] = pressed;

    if (m_callbacks.onMouse)
        m_callbacks.onMouse(button, pressed);
}

void Window::HandleScrollMessage(float offset)
{
    if (m_callbacks.onScroll)
        m_callbacks.onScroll(offset);
}

LRESULT CALLBACK Window::WindowProc(HWND hWnd, UINT msg, WPARAM wp, LPARAM lp)
{
    Window* self = nullptr;

    if (msg == WM_NCCREATE)
    {
        CREATESTRUCTW* cs = reinterpret_cast<CREATESTRUCTW*>(lp);
        self = reinterpret_cast<Window*>(cs->lpCreateParams);

        SetWindowLongPtrW(hWnd, GWLP_USERDATA, (LONG_PTR)self);
        return TRUE;
    }

    self = reinterpret_cast<Window*>(GetWindowLongPtrW(hWnd, GWLP_USERDATA));

    if (!self)
        return DefWindowProcW(hWnd, msg, wp, lp);

    switch (msg)
    {
        case WM_CLOSE:
            self->m_shouldClose = true;
            return 0;

        case WM_KEYDOWN:
        case WM_SYSKEYDOWN:
            self->HandleKeyMessage(wp, true);
            return 0;

        case WM_KEYUP:
        case WM_SYSKEYUP:
            self->HandleKeyMessage(wp, false);
            return 0;

        case WM_LBUTTONDOWN:
            self->HandleMouseButtonMessage(MouseButton::Left, true);
            return 0;
        case WM_LBUTTONUP:
            self->HandleMouseButtonMessage(MouseButton::Left, false);
            return 0;

        case WM_RBUTTONDOWN:
            self->HandleMouseButtonMessage(MouseButton::Right, true);
            return 0;
        case WM_RBUTTONUP:
            self->HandleMouseButtonMessage(MouseButton::Right, false);
            return 0;

        case WM_MOUSEWHEEL:
        {
            float offset = GET_WHEEL_DELTA_WPARAM(wp) / 120.0f;
            self->HandleScrollMessage(offset);
            return 0;
        }
    }

    return DefWindowProcW(hWnd, msg, wp, lp);
}

void Window::SetCallbacks(const WindowCallbacks &callbacks)
{
    m_callbacks = callbacks;
}

void Window::DrawBuffer(Framebuffer &framebuffer)
{
    if (!m_handle || !m_memoryDC || !framebuffer.GetColorBuffer())
        return;

    unsigned char* dst = m_surface->GetLdrPixel(0, 0);
    unsigned char* src = framebuffer.GetColorBuffer();

    int totalBytes = framebuffer.GetWidth() * framebuffer.GetHeight() * 4;
    std::copy(src, src + totalBytes, dst);

    HDC windowDC = GetDC(m_handle);
    BitBlt(
        windowDC,
        0, 0,
        framebuffer.GetWidth(),
        framebuffer.GetHeight(),
        m_memoryDC,
        0, 0,
        SRCCOPY
    );
    ReleaseDC(m_handle, windowDC);
}

float Window::GetTime()
{
    static double frequency = 0.0;
    static double initialTime = 0.0;

    LARGE_INTEGER counter;
    LARGE_INTEGER freq;

    if (frequency == 0.0)
    {
        QueryPerformanceFrequency(&freq);
        frequency = double(freq.QuadPart);

        QueryPerformanceCounter(&counter);
        initialTime = double(counter.QuadPart) / frequency;
    }

    QueryPerformanceCounter(&counter);
    return float(double(counter.QuadPart) / frequency - initialTime);
}

void Window::CreateSurface(int width, int height)
{
    assert(width > 0 && height > 0);

    if (m_surface)
        delete m_surface;

    m_surface = new Image(width, height, 4, ImageFormat::LDR);

    HDC windowDC = GetDC(m_handle);

    if (m_memoryDC)
    {
        DeleteDC(m_memoryDC);
        m_memoryDC = nullptr;
    }

    m_memoryDC = CreateCompatibleDC(windowDC);
    assert(m_memoryDC != nullptr);

    BITMAPINFOHEADER biHeader{};
    biHeader.biSize = sizeof(BITMAPINFOHEADER);
    biHeader.biWidth = width;
    biHeader.biHeight = -height;  
    biHeader.biPlanes = 1;
    biHeader.biBitCount = 32;     
    biHeader.biCompression = BI_RGB;

    void* pixels = nullptr;
    HBITMAP hBitmap = CreateDIBSection(
        m_memoryDC,
        reinterpret_cast<BITMAPINFO*>(&biHeader),
        DIB_RGB_COLORS,
        &pixels,
        nullptr,
        0
    );
    assert(hBitmap != nullptr);

    HBITMAP oldBitmap = (HBITMAP)SelectObject(m_memoryDC, hBitmap);
    if (oldBitmap)
        DeleteObject(oldBitmap);

    m_surface->SetLdrBuffer(reinterpret_cast<unsigned char*>(pixels));
    m_surface->SetOwnsMemory(false);  

    ReleaseDC(m_handle, windowDC);
}

void Window::PresentSurface()
{
    if (!m_handle || !m_memoryDC)
        return;

    HDC windowDC = GetDC(m_handle);

    int width  = m_surface->GetWidth();
    int height = m_surface->GetHeight();

    BitBlt(
        windowDC,
        0, 0,
        width,
        height,
        m_memoryDC,
        0, 0,
        SRCCOPY
    );

    ReleaseDC(m_handle, windowDC);
}

